import torch.nn as nn
from get_dict import get_data
from dataset import MyPairDataset
from torch.utils.data import DataLoader
from encoder import MyEncoderGRU

# 调用获取数据函数
english_word2index, english_index2word, english_word_n, \
    french_word2index, french_index2word, french_word_n, my_pairs = get_data()

# GRU解码器
class MyDecoderGRU(nn.Module):
    def __init__(self, vocab_size, hidden_size):
        super(MyDecoderGRU, self).__init__()
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size

        # 定义词嵌入层
        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=hidden_size)

        self.gru = nn.GRU(input_size=hidden_size, hidden_size=hidden_size, num_layers=1, batch_first=True)

        self.linear = nn.Linear(in_features=hidden_size, out_features=vocab_size)

        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self,input, hidden):
        embed = self.embedding(input)

        output, hidden = self.gru(embed, hidden)

        output = self.linear(output[-1])

        result = self.softmax(output)

        return result, hidden

if __name__ == '__main__':
    my_dataset = MyPairDataset(my_pairs)

    dataloader = DataLoader(dataset=my_dataset, batch_size=1, shuffle=True)

    for x, y in dataloader:
        print('目标y的形状:', y.shape)

        # 获取编码器最后一个时间步隐藏层输出作为解码器的第一个时间步的隐藏层输入
        my_encoder_gru = MyEncoderGRU(vocab_size=english_word_n, hidden_size=10)
        hidden = my_encoder_gru.init_hidden()
        encoder_output, hidden = my_encoder_gru(x, hidden)

        my_decoder = MyDecoderGRU(vocab_size=french_word_n, hidden_size=10)
        print('解码器模型结构:', my_decoder)
        for i in range(y.shape[1]):
            print(f"========解码第{i + 1}次========")
            input_y = y[0][i]
            print(f'输入的y的形状：-->{input_y.shape}')
            input_y = input_y.view(1, -1)
            print(f'输入的y的形状：-->{input_y.shape}')
            output, hidden = my_decoder(input_y, hidden)
            print(f'解码器输出的形状：-->{output.shape}')
            print(f'解码器输出的结果：-->{output}')
            print(f'解码器hidden的形状：-->{hidden.shape}')
            print(f'解码器hidden的结果：-->{hidden}')
        print('=============y的形状变换=============')
        print('经过dataset', my_dataset[0][1].shape)
        print('经过dataloader', y.shape)
        print('经过编码器模型', output.shape)
        break